import os
import json
import requests
from typing import Optional
from dotenv import load_dotenv
import pymysql
from langchain_deepseek import ChatDeepSeek

# 加载环境变量
load_dotenv(override=True)

class LLMText2SQL:
    """使用本地LLM实例进行自然语言到SQL的转换"""
    
    def __init__(self):
        # MySQL配置
        self.mysql_config = {
            'host': os.getenv('HOST', 'localhost'),
            'port': int(os.getenv('PORT', 3306)),
            'user': os.getenv('USER', 'root'),
            'password': os.getenv('MYSQL_PW', ''),
            'database': os.getenv('DB_NAME', 'user_insight'),
        }
        self.table_name = os.getenv('TABLE_NAME', 'insight_general_table')
        
        # 初始化本地LLM实例
        self._llm_model_name = os.getenv("LLM_MODEL", "deepseek-chat")
        try:
            self._llm = ChatDeepSeek(model=self._llm_model_name)
            print(f"✅ 统算模式LLM已配置: {self._llm_model_name}")
        except Exception as e:
            print(f"⚠️ LLM配置失败: {e}")
            self._llm = None
    
    def get_table_schema(self) -> str:
        """获取数据库表的schema信息，用于LLM生成SQL的prompt"""
        connection = self._connect_mysql()
        if not connection:
            return "无法连接数据库"
        
        try:
            with connection.cursor() as cursor:
                # 获取表结构
                cursor.execute(f"DESCRIBE {self.mysql_config['database']}.{self.table_name}")
                columns = cursor.fetchall()
                
                # 获取行数
                cursor.execute(f"SELECT COUNT(*) FROM {self.mysql_config['database']}.{self.table_name}")
                row_count = cursor.fetchone()[0]
                
                # 获取示例数据
                cursor.execute(f"SELECT * FROM {self.mysql_config['database']}.{self.table_name} LIMIT 3")
                sample_data = cursor.fetchall()
                
                # 获取列名
                cursor.execute(f"SHOW COLUMNS FROM {self.mysql_config['database']}.{self.table_name}")
                column_names = [col[0] for col in cursor.fetchall()]
                
                # 构建schema信息
                schema_info = f"表名: {self.mysql_config['database']}.{self.table_name}\n"
                schema_info += f"总行数: {row_count}\n\n"
                schema_info += "字段结构:\n"
                
                for col in columns:
                    field_name = col[0]
                    field_type = col[1]
                    nullable = "NULL" if col[2] == "YES" else "NOT NULL"
                    default_value = col[4] if col[4] else "无默认值"
                    schema_info += f"- {field_name}: {field_type} ({nullable}), 默认值: {default_value}\n"
                
                schema_info += "\n示例数据:\n"
                if sample_data:
                    for i, row in enumerate(sample_data):
                        schema_info += f"行{i+1}: "
                        for j, value in enumerate(row):
                            schema_info += f"{column_names[j]}={value}, "
                        schema_info = schema_info.rstrip(", ") + "\n"
                
                return schema_info
                
        except Exception as e:
            return f"获取schema失败: {e}"
        finally:
            connection.close()
    
    def text_to_sql_with_llm(self, user_request: str) -> str:
        """使用大模型根据用户需求生成SQL查询"""
        try:
            # 获取表结构信息
            schema_info = self.get_table_schema()
            
            # 构建LLM prompt
            prompt = f"""
你是一个SQL专家。请根据用户的自然语言需求，生成对应的MySQL查询语句。

数据库表信息：
{schema_info}

用户需求：{user_request}

要求：
1. 只返回SQL语句，不要其他解释文字
2. 使用标准MySQL语法
3. 字段名请使用表中实际存在的字段
4. 如果用户询问省份相关信息，请从用户需求中提取具体省份名称（如北京、山东等）
5. 如果用户询问年龄分布，请按ATTR_AGE字段分组，将NULL值处理为'未知'
6. 如果用户询问ARPU相关，请使用METRIC_ARPU字段
7. 请根据用户需求智能选择GROUP BY的维度和聚合函数
8. 确保SQL查询结果包含合适的列名，便于图表展示

示例SQL格式：
SELECT 列名1 AS '显示名1', COUNT(DISTINCT USER_ID) AS '用户数量' 
FROM {self.mysql_config['database']}.{self.table_name} 
WHERE 条件 
GROUP BY 列名1

请生成SQL：
"""
            
            # 调用本地LLM实例
            if self._llm is not None:
                try:
                    resp = self._llm.invoke(prompt)
                    content = getattr(resp, "content", resp)
                    if not isinstance(content, str):
                        content = str(content)
                    
                    # 清理SQL，移除markdown标记
                    import re
                    sql_query = re.sub(r'```sql\s*', '', content)
                    sql_query = re.sub(r'```\s*', '', sql_query)
                    sql_query = sql_query.strip()
                    
                    # 验证SQL基本格式
                    if sql_query and 'SELECT' in sql_query.upper():
                        print(f"✅ LLM生成的SQL查询: {sql_query}")
                        return sql_query
                    else:
                        print(f"⚠️ LLM生成的SQL格式不正确: {sql_query}")
                        return self._generate_fallback_sql(user_request)
                        
                except Exception as e:
                    print(f"⚠️ LLM调用失败: {e}")
                    return self._generate_fallback_sql(user_request)
            else:
                print(f"⚠️ LLM未配置，使用备用逻辑")
                return self._generate_fallback_sql(user_request)
                
        except Exception as e:
            print(f"SQL生成失败: {e}")
            return self._get_default_sql_query()
    
    def _generate_fallback_sql(self, user_request: str) -> str:
        """当LLM不可用时的备用SQL生成逻辑"""
        try:
            user_request_lower = user_request.lower()
            
            # 智能提取省份信息
            provinces = []
            province_keywords = ['北京', '山东', '云南', '上海', '广东', '浙江', '江苏', '河北', '河南', '湖北', '湖南', '四川', '重庆', '天津', '辽宁', '吉林', '黑龙江', '内蒙古', '山西', '陕西', '甘肃', '青海', '宁夏', '新疆', '西藏', '广西', '海南', '贵州', '福建', '江西', '安徽']
            
            for province in province_keywords:
                if province in user_request:
                    provinces.append(f"'{province}'")
            
            # 如果没有指定省份，使用默认值（只用北京）
            if not provinces:
                provinces = ["'北京'"]
            
            provinces_str = ", ".join(provinces)
            print(f"🔍 备用逻辑提取的省份: {provinces}")
            
            # 提取ARPU条件
            arpu_condition = ""
            if 'arpu' in user_request_lower:
                import re
                numbers = re.findall(r'\d+', user_request)
                if numbers and ('大于' in user_request or '>' in user_request):
                    arpu_condition = f"AND METRIC_ARPU > {numbers[0]}"
                elif numbers and ('小于' in user_request or '<' in user_request):
                    arpu_condition = f"AND METRIC_ARPU < {numbers[0]}"

            # 检查是否询问年龄分布
            ask_age_distribution = any(k in user_request for k in ["年龄分布", "按年龄", "年龄段", "年龄结构", "各年龄", "年龄占比", "年龄统计"]) or ('年龄' in user_request and ('分布' in user_request or '占比' in user_request or '统计' in user_request))

            if ask_age_distribution:
                # 年龄分布查询
                sql_query = f"""
                SELECT 
                    CASE WHEN ATTR_AGE IS NULL THEN '未知' ELSE CAST(ATTR_AGE AS CHAR) END AS '年龄',
                    COUNT(DISTINCT USER_ID) AS '用户数量'
                FROM {self.mysql_config['database']}.{self.table_name}
                WHERE PROV_ID IN ({provinces_str}) {arpu_condition}
                GROUP BY 1
                ORDER BY 
                    CASE WHEN ATTR_AGE IS NULL THEN 1 ELSE 0 END ASC,
                    CASE WHEN ATTR_AGE IS NULL THEN 0 ELSE ATTR_AGE END ASC
                """
            else:
                # 默认按省份聚合
                sql_query = f"""
                SELECT 
                    PROV_ID AS '省份', 
                    COUNT(DISTINCT USER_ID) AS '用户数量' 
                FROM {self.mysql_config['database']}.{self.table_name} 
                WHERE PROV_ID IN ({provinces_str}) {arpu_condition}
                GROUP BY PROV_ID
                """
            
            print(f"✅ 备用逻辑生成的SQL查询: {sql_query}")
            return sql_query.strip()
                
        except Exception as e:
            print(f"备用SQL生成失败: {e}")
            return self._get_default_sql_query()
    
    def _get_default_sql_query(self) -> str:
        """获取默认的SQL查询（当LLM不可用时）"""
        return f"""
        SELECT 
            PROV_ID AS '省份', 
            COUNT(DISTINCT USER_ID) AS '用户数量' 
        FROM {self.mysql_config['database']}.{self.table_name} 
        WHERE PROV_ID IN ('北京') 
        GROUP BY PROV_ID
        """
    
    def _connect_mysql(self):
        """连接MySQL数据库"""
        try:
            connection = pymysql.connect(**self.mysql_config)
            return connection
        except Exception as e:
            print(f"MySQL 连接失败: {e}")
            return None
    
    def execute_sql_query(self, sql_query: str) -> Optional[list]:
        """执行SQL查询并返回结果"""
        connection = self._connect_mysql()
        if not connection:
            return None
        
        try:
            with connection.cursor() as cursor:
                cursor.execute(sql_query)
                results = cursor.fetchall()
                
                # 获取查询结果的列名（使用cursor.description）
                if cursor.description:
                    columns = [desc[0] for desc in cursor.description]
                else:
                    # 如果没有description，使用默认列名
                    columns = [f"column_{i}" for i in range(len(results[0]) if results else 0)]
                
                # 将结果转换为字典列表
                formatted_results = []
                for row in results:
                    row_dict = {}
                    for i, value in enumerate(row):
                        if i < len(columns):
                            row_dict[columns[i]] = value
                    formatted_results.append(row_dict)
                
                print(f"✅ SQL查询执行成功，返回 {len(formatted_results)} 条记录")
                print(f"📊 查询结果列名: {columns}")
                print(f"📋 查询结果示例: {formatted_results[:2] if formatted_results else '无数据'}")
                
                return formatted_results
                
        except Exception as e:
            print(f"SQL执行失败: {e}")
            return None
        finally:
            connection.close()
    
    def get_llm_status(self) -> str:
        """获取LLM配置状态"""
        if self._llm is not None:
            return f"✅ 本地LLM服务可用 (模型: {self._llm_model_name})"
        else:
            return f"❌ 本地LLM服务不可用"
